// 演示 LLMChainExtractor 对文档进行压缩
import { ChatOllama } from "@langchain/ollama";
import { Document } from "langchain/document";
import { LLMChainExtractor } from "langchain/retrievers/document_compressors/chain_extract";

const llm = new ChatOllama({
  model: "llama3",
  temperature: 0.7,
});

// 创建了一个压缩器
const extractor = LLMChainExtractor.fromLLM(llm);

// 构建一段文档
const docs = [
  new Document({
    pageContent: `孔乙己走到酒店，说“要一碟茴香豆”。茴香豆常作为下酒小菜，价格便宜，常见于短衣帮和穿长衫的顾客点酒时。`,
  }),
  new Document({
    pageContent: `另一个段落：今天的天气很好，阳光明媚。`,
  }),
  new Document({
    pageContent: `游泳池边上，游过一群小黄鸭`,
  }),
  new Document({
    pageContent: `茴香豆可以拿来泡茶，上次我就看到有一个大婶儿拿茴香豆来泡茶`,
  }),
];

const query = "茴香豆的作用";

const res = await extractor.compressDocuments(docs, query);

console.log(res);
